#include "twofish.h"
#include "types.h"
#include "h_func.h"
#include "matrix.h"
#include "tables.h"
#include "mul_tab.h"
#include "field.h"
#include "pht.h"
#include "rotate.h"
#include "set_error_insert.h"

#include <string.h>

void input_whitening(byte* plaintext, key_type* encrypt_key)
{
	int i,j;
	byte tmp_bytes[4];
	for( i = 0; i < 4; ++i)
	{
		split_word(encrypt_key->expanded_key[i], tmp_bytes, 4);
		for( j = 0; j < 4; ++j)
		{
			plaintext[i * 4 + j] ^= tmp_bytes[j];
		}
	}
}

void output_whitening(byte* mid_data, key_type* encrypt_key)
{
	int i,j;
	byte tmp_bytes[4];
	for( i = 0; i < 4; ++i)
	{
		split_word(encrypt_key->expanded_key[4+i], tmp_bytes, 4);
		for( j = 0; j < 4; ++j)
		{
			mid_data[i * 4 + j] ^= tmp_bytes[j];
		}
	}
}

extern byte gf_256_169_mul(byte lhs, byte rhs);

word g_func_word(word value, key_type* encrypt_key)
{
	value = h_func(value, encrypt_key->s_array, encrypt_key->key_size / 8);
	return value;
//	byte tmp_bytes[4], tmp_bytes2[4];
//	split_word(value, tmp_bytes, 4);
//	byte_matrix_mul(mds, 4, 4, tmp_bytes, 4, 1, tmp_bytes2, 4, element_add, gf_256_169_mul);

//	return join_word(tmp_bytes2);
}
//values must be 4 bytes
word g_func_byte(byte* values, key_type* encrypt_key)
{
	word tmp = join_word(values);
	return g_func_word(tmp, encrypt_key);
}

//#define LAST_ROUND_OUTPUT

#include <stdio.h>
void round_func(byte* left_data, byte* right_data, byte* result, key_type* encrypt_key, int round)
{
#ifdef LAST_ROUND_OUTPUT
	if(round == 15)
	{
		word tmp1 = join_word(left_data);
		word tmp2 = join_word(left_data+4);
		printf("Before sbox,\t\t left0:%.8x, left1:%.8x\n", tmp1, tmp2);
	}
#endif

	word left0 = g_func_byte(left_data, encrypt_key);
	word left1 = join_word(left_data + 4);
	left1 = rol(left1, 8);
	left1 = g_func_word(left1, encrypt_key);

#ifdef LAST_ROUND_OUTPUT
	if(round == 15)
		printf("After g function,\t left0:%.8x, left1:%.8x\n", left0, left1);
#endif

	pht(left0, left1, &left0, &left1);

#ifdef LAST_ROUND_OUTPUT
	if(round == 15)
		printf("After pht,\t\t left0:%.8x, left1:%.8x\n", left0, left1);
#endif

	left0 = (left0 + encrypt_key->expanded_key[2*round+8]);
	left1 = (left1 + encrypt_key->expanded_key[2*round+9]);

#ifdef LAST_ROUND_OUTPUT
	if(round == 15)
		printf("After add round key,\t left0:%.8x, left1:%.8x\n", left0, left1);
#endif

	word right0 = join_word(right_data);
	word right1 = join_word(right_data + 4);

	left0 = ror(left0 ^ right0, 1);
	left1 = rol(right1, 1) ^ left1;

#ifdef LAST_ROUND_OUTPUT
	if(round == 15)
		printf("After rotate and xor,\t left0:%.8x, left1:%.8x\n", left0, left1);
#endif
	
	split_word( left0, result, 4);
	split_word( left1, result + 4, 4);

}

#define ROUND_PRINT
#ifdef ROUND_PRINT
static int round_print_enabled = 0;
static int round_print_round = 0;

void enable_round_print(int enabled)
{
	round_print_enabled = enabled;
}
void set_round_print(int round)
{
	round_print_round = round;
}
#endif

//plaintext's size must be 128bits
void twofish_128_encrypt(byte* plaintext, byte* cipher, key_type* encrypt_key)
{
	memcpy(cipher, plaintext, BLOCK_SIZE);
	input_whitening(cipher, encrypt_key);

	byte *left_part, *right_part;
	byte tmp[BLOCK_SIZE/2];
	int round;
	for( round = 0; round < ROUND; ++round)
	{
		byte* des;
		if( round % 2 == 0)
		{
			left_part = cipher;
			right_part = cipher + BLOCK_SIZE / 2;
		}
		else
		{
			left_part = cipher + BLOCK_SIZE /2 ;
			right_part = cipher;
		}

		//call error_insert hook function
		error_insert(left_part, right_part, round);

		des = right_part;
		round_func(left_part, right_part, des, encrypt_key, round);

#ifdef ROUND_PRINT
		if(round_print_enabled)
		{
			if(round_print_round == round)
			{
				int i = 0;
				//after swap right_part is left_part
				for(i = 0; i < BLOCK_SIZE/2; ++i)
				{
					printf("%.2x ", right_part[i]);
				}
				for(i = 0; i < BLOCK_SIZE/2; ++i)
				{
					printf("%.2x ", left_part[i]);
				}
				printf("\n");
			}
		}
#endif
	}

	//undo last swap
	memcpy(tmp, cipher, BLOCK_SIZE /2);
	memcpy(cipher, cipher + BLOCK_SIZE/2 ,BLOCK_SIZE/2);
	memcpy(cipher + BLOCK_SIZE/2 ,tmp, BLOCK_SIZE/2);

	output_whitening(cipher, encrypt_key);

}

#include "expandkey.h"
#include <stdio.h>
#include <assert.h>
void test_twofish_128_encrypt()
{
	printf("Start testing twofish_128_encrypt.............\n");


	const char* test_file_names[] = {
		"test_encrypt_128.txt",
		"test_encrypt_192.txt",
		"test_encrypt_256.txt"
	};
	int i,j,k;
	byte key[MAX_KEY_LEN] = {
		0
	};
	byte plaintext[BLOCK_SIZE] = {
		0
	};
	byte cipher[BLOCK_SIZE] = {
		0
	};
	byte real_cipher[BLOCK_SIZE] = {
		0
	};
	int key_size;
	key_type encrypt_key;
	
	const int buffer_size = 1024;
	char buffer[buffer_size];
	FILE* test_file;
	for( i = 0; i < sizeof(test_file_names) /sizeof(test_file_names[0]); ++i)
	{
		test_file = fopen(test_file_names[i], "rb");
		if(test_file == NULL)
		{
			printf("Can't open %s for read!\n", test_file_names[i]);
			continue;
		}

		//get key size and key
		unsigned int tmp;
		fscanf(test_file, "%d", &key_size);
		key_size /= 8;

		fscanf(test_file, "%s", buffer);
		for( j = 0; j < key_size; ++j)
		{
			//use tmp to avoid not enough space bug
			char tmp_char = buffer[2*(j+1)];
			buffer[2*(j+1)]= 0;
			sscanf(buffer + 2*j, "%x", &tmp);
			key[j] = tmp;
			buffer[2*(j+1)] = tmp_char;
		}
		pre_expand_key(key, key_size, &encrypt_key);

		int example_count;
		fscanf(test_file, "%d", &example_count);
		for( j = 0; j < example_count; ++j)
		{
			fscanf(test_file, "%s", buffer);
			for( k = 0;k < BLOCK_SIZE; ++k)
			{
				char tmp_char = buffer[2*(k+1)];
				buffer[2*(k+1)]= 0;
				sscanf(buffer + 2*k, "%x", &tmp);
				plaintext[k] = tmp;
				buffer[2*(k+1)] = tmp_char;
			}
			fscanf(test_file, "%s", buffer);
			for( k = 0;k < BLOCK_SIZE; ++k)
			{
				char tmp_char = buffer[2*(k+1)];
				buffer[2*(k+1)]= 0;
				sscanf(buffer + 2*k, "%x", &tmp);
				real_cipher[k] = tmp;
				buffer[2*(k+1)] = tmp_char;
			}
			twofish_128_encrypt(plaintext, cipher, &encrypt_key);

			for( k = 0; k < BLOCK_SIZE; ++k)
			{
				assert(real_cipher[k] == cipher[k]);
			}
		}

		fclose(test_file);
	}

	
//	byte plaintext[][BLOCK_SIZE] = {
//		{0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00,},
//		{0x80,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00,},
//		{0x40,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00,},
//	};
//	byte real_cipher[][BLOCK_SIZE] = {
//		{0x9F,0x58,0x9F,0x5C,0xF6,0x12,0x2C,0x32,0xB6,0xBF,0xEC,0x2F,0x2A,0xE8,0xC3,0x5A,},
//		{0x73,0xB9,0xFF,0x14,0xCF,0x25,0x89,0x90,0x1F,0xF5,0x2A,0x0D,0x6F,0x4B,0x7E,0xDE,},
//		{0xF5,0xA9,0x15,0x0B,0xAB,0x6D,0x6A,0xEB,0xD6,0xB4,0xF9,0x7D,0x9E,0x93,0xB2,0x8B,},
//	};
//
//	int plaintext_size = sizeof(plaintext) / sizeof(plaintext[0]);
//
//	byte cipher[BLOCK_SIZE] = {
//		0
//	};
//	
//
////	for( i = 0; i < 2; ++i)
////	{
////		printf("%x\n", encrypt_key.s_array[i]);
////	}
//
//
//	for(i = 0; i < plaintext_size; ++i)
//	{
//		twofish_128_encrypt(plaintext[i], cipher, &encrypt_key);
//		for(j = 0; j < BLOCK_SIZE; ++j)
//		{
//			assert(cipher[j] == real_cipher[i][j]);
//		}
//	}
//
	printf("Test twofish_128_encrypt complete!\n");
}
